# -*- coding: utf-8 -*-
import json

import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from sparkai.core.messages import ChatMessage
from sparkai.llm.llm import ChatSparkLLM

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")


############### 1. 加载api ###############
SPARKAI_V4_API_KEY = "7570a4c16f83a321333d3fde1c28ec00"
SPARKAI_V4_API_SECRET = "MWU4MmQ4N2I2MDExYjU5YzI4ODk2YTY4"
SPARKAI_V4_APP_ID = "f542483a"
SPARKAI_V4_URL = "wss://spark-api.xf-yun.com/v4.0/chat"
SPARKAI_V4_DOMAIN = "4.0Ultra"

spark = ChatSparkLLM(
    spark_api_url=SPARKAI_V4_URL,
    spark_app_id=SPARKAI_V4_APP_ID,
    spark_api_key=SPARKAI_V4_API_KEY,
    spark_api_secret=SPARKAI_V4_API_SECRET,
    spark_llm_domain=SPARKAI_V4_DOMAIN,
    streaming=True,
)


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})


@app.get("/change")
async def change_key(key: str):
    return "Can't change sparkai key"


@app.post("/chat")
async def chat(request: Request):
    form = await request.form()
    messages = form["prompts"]
    # api_key = form.get("apiKey")

    if not messages:
        return JSONResponse(
            status_code=400,
            content={
                "error": {
                    "message": "请输入prompts！",
                    "type": "invalid_request_error",
                    "code": "",
                }
            },
        )

    prompts = json.loads(messages)
    messages = [
        ChatMessage(
            role=item["role"],
            content=item["content"],
        )
        for item in prompts
    ]

    def generate():
        response = spark.stream(messages, temperature=0.7)
        for message in response:
            if not message.content:
                continue
            yield message.content

    return "".join(generate())


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8899)
